import arrow
import math
import time
import threading

from app.db.models.tasks import Task
from app.db.models.wass import ScanEvents
from app.celery.handle_result.utils import save_spider_urls
from app.config.settings import SPIDER_URL_TIME_STR


def update_data(thread_num, task_id_list, task_flag_dict):
    for _ in task_id_list:
        task_id, task_session_id = _
        update_dict = {}
        if scan_count := ScanEvents._get_db()["scan_pages"].find({"task_session_id": task_session_id}).count():
            update_dict.update({"result.addition.statistics.scan_count": scan_count})
        if url_count := save_spider_urls(task_id):
            update_dict.update({"result.addition.statistics.urls": url_count})
        if update_dict:
            Task._get_collection().update({"_id": task_id}, {"$set": update_dict})
        task_flag_dict.update({task_id: 1})
    print(f"线程{thread_num}结束， 处理了{len(task_id_list)}个")


def main():
    task_list = [
            (t.id, t.taskSessionId) for t in Task.objects.filter(
                taskType__in=['vul', 'ssl', 'securityEvent', 'content'],
                status="completed",
                endTime__gte=arrow.get(SPIDER_URL_TIME_STR).datetime
            ).only("id", "taskSessionId")
        ]
    count = len(task_list)
    task_flag_dict = {t[0]: 0 for t in task_list}
    t_count = 20
    limit = math.ceil(count / t_count)
    print("开始更新......")
    for thread_num in range(0, t_count):
        _ = task_list[thread_num * limit: (thread_num+1) * limit]
        thread = threading.Thread(target=update_data, args=(thread_num, _, task_flag_dict))
        thread.start()

    while sum(task_flag_dict.values()) < count:
        print(f"总共的需要处理的任务数{count}，已经处理 {sum(task_flag_dict.values())}个")
        time.sleep(5)
    print(f"总共的需要处理的任务数{count}，已经处理 {sum(task_flag_dict.values())}个")
    print("结束")



